{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# rStar-Math: Geometry and Proofs\n",
    "\n",
    "This notebook demonstrates how rStar-Math handles geometric problems and mathematical proofs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.patches import Circle, Rectangle, Polygon\n",
    "from src.core.mcts import MCTS\n",
    "from src.core.ppm import ProcessPreferenceModel\n",
    "from src.models.model_interface import ModelFactory"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup Components"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "mcts = MCTS.from_config_file('config/default.json')\n",
    "ppm = ProcessPreferenceModel.from_config_file('config/default.json')\n",
    "model = ModelFactory.create_model('openai', os.getenv('OPENAI_API_KEY'), 'config/default.json')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Basic Geometry Problems"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "geometry_problems = [\n",
    "    \"Find the area of a circle with radius 5\",\n",
    "    \"Calculate the volume of a sphere with radius 3\",\n",
    "    \"In a right triangle with legs 3 and 4, find the hypotenuse\",\n",
    "    \"Find the area of a regular hexagon with side length 2\"\n",
    "]\n",
    "\n",
    "for problem in geometry_problems:\n",
    "    print(f\"Problem: {problem}\\n\")\n",
    "    action, trajectory = mcts.search(problem)\n",
    "    \n",
    "    print(\"Solution Steps:\")\n",
    "    for step in trajectory:\n",
    "        confidence = ppm.evaluate_step(step['state'], model)\n",
    "        print(f\"- {step['state']}\")\n",
    "        print(f\"  Confidence: {confidence:.2f}\\n\")\n",
    "    print(\"-\" * 50 + \"\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Geometric Proofs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "proof_problems = [\n",
    "    \"Prove that the sum of angles in a triangle is 180 degrees\",\n",
    "    \"Prove the Pythagorean theorem\",\n",
    "    \"Prove that the diagonals of a rectangle are equal\"\n",
    "]\n",
    "\n",
    "for problem in proof_problems:\n",
    "    print(f\"Problem: {problem}\\n\")\n",
    "    action, trajectory = mcts.search(problem)\n",
    "    \n",
    "    print(\"Proof Steps:\")\n",
    "    for step in trajectory:\n",
    "        confidence = ppm.evaluate_step(step['state'], model)\n",
    "        print(f\"- {step['state']}\")\n",
    "        print(f\"  Confidence: {confidence:.2f}\\n\")\n",
    "    print(\"-\" * 50 + \"\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Visualization of Geometric Concepts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "source": [
    "def plot_right_triangle():\n",
    "    \"\"\"Visualize Pythagorean theorem.\"\"\"\n",
    "    fig, ax = plt.subplots(figsize=(8, 8))\n",
    "    \n",
    "    # Draw triangle\n",
    "    ax.plot([0, 3, 0, 0], [0, 0, 4, 0], 'b-', linewidth=2)\n",
    "    \n",
    "    # Add squares on each side\n",
    "    ax.add_patch(Rectangle((0, 0), 3, 3, fill=False, color='red'))\n",
    "    ax.add_patch(Rectangle((3, 0), 4, 4, fill=False, color='green', angle=90))\n",
    "    ax.add_patch(Polygon([[0, 0], [0, 4], [-5, 4], [-5, 0]], fill=False, color='blue'))\n",
    "    \n",
    "    # Add labels\n",
    "    ax.text(1.5, -0.5, '3', ha='center')\n",
    "    ax.text(-0.5, 2, '4', va='center')\n",
    "    ax.text(1, 2, '5', ha='center', va='center', rotation=-45)\n",
    "    \n",
    "    ax.set_aspect('equal')\n",
    "    ax.grid(True)\n",
    "    plt.title(\"Pythagorean Theorem: a² + b² = c²\")\n",
    "    plt.show()\n",
    "\n",
    "def plot_circle_properties():\n",
    "    \"\"\"Visualize circle properties.\"\"\"\n",
    "    fig, ax = plt.subplots(figsize=(8, 8))\n",
    "    \n",
    "    # Draw circle\n",
    "    circle = Circle((0, 0), 5, fill=False)\n",
    "    ax.add_patch(circle)\n",
    "    \n",
    "    # Draw radius\n",
    "    ax.plot([0, 5], [0, 0], 'r-', label='Radius')\n",
    "    \n",
    "    # Draw diameter\n",
    "    ax.plot([-5, 5], [0, 0], 'g--', label='Diameter')\n",
    "    \n",
    "    ax.set_aspect('equal')\n",
    "    ax.grid(True)\n",
    "    plt.legend()\n",
    "    plt.title(\"Circle Properties\")\n",
    "    plt.show()\n",
    "\n",
    "# Generate visualizations\n",
    "plot_right_triangle()\n",
    "plot_circle_properties()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
